package edu.gatech.cc.liam.marl.decpomdp;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import edu.gatech.cc.liam.core.rl.models.CartesianIterator;
import edu.gatech.cc.liam.core.rl.models.JointIterator;
import edu.gatech.cc.liam.core.rl.models.decpomdp.DecMDP;
import edu.gatech.cc.liam.core.rl.models.decpomdp.DecPOMDP;
import edu.gatech.cc.liam.core.rl.models.decpomdp.DecPOMDPFactory;
import edu.gatech.cc.liam.geometry.linear.NPoint;

public class PPL0Approximator {
	
	DecPOMDP problem; // base problem trying to solve
	List<Belief>[] l0Beliefs;  // the level 0 beliefs over which agents will reason and form nested beliefs
	DecMDP approxP; // the approximation of the problem
	
	public PPL0Approximator(DecPOMDP problem, List<Belief>[] level0Beliefs) {
		this.problem = problem;
		this.l0Beliefs = level0Beliefs;
		makel0BeliefApprox();
	}
	
	private void makel0BeliefApprox() {
		int numJActions = problem.numJActions;
		// construct types
		@SuppressWarnings("unchecked")
		ArrayList<String>[] types = new ArrayList[problem.numAgents + 1];  // an extra agent for nature
		types[0] = new ArrayList<String>();
		for(int s=0; s<problem.numStates; s++) {
			types[0].add(problem.states.get(s));
		}
		for(int a = 0; a<problem.numAgents; a++) {
			types[a+1] = new ArrayList<String>();
			for(Belief b : l0Beliefs[a])
				types[a+1].add(b.toNPointString());
		}
		// a "noop" for nature is added to actions 
		@SuppressWarnings("unchecked")
		List<String>[] actions = new List[problem.numAgents + 1];
		actions[0] = Arrays.asList(new String[]{"noop"});
		for(int i=0; i < problem.numAgents; i++)
			actions[i+1] = problem.actions[i];
		// discount is the same
		approxP = new DecMDP(types, actions, problem.discount);

		// construct transition function
		List<List<String>> allJointTypes = 	DecMDP.makeJointList(types, approxP.typeSS);
		int numJTypes = allJointTypes.size();
		double[][][] transitionFunction = 
			new double[numJTypes][numJActions][numJTypes]; // [start-type][actions][end-type]
		for(int[] jStartT : new CartesianIterator(approxP.typeSS)) {
			int jStartTi = JointIterator.getJointIndex(approxP.typeSS, jStartT);
			for(int jAi = 0; jAi < numJActions; jAi++) {
				// for each starting joint-type and joint-action
				// sum over successor true states and observations (each yielding a successor type)
				for(int nextTrueState = 0; nextTrueState < problem.numStates; nextTrueState++) {
					for(int[] jO : new CartesianIterator(problem.obsSS)) {
						int jOi = JointIterator.getJointIndex(problem.obsSS, jO);
						int nextTi = getNextJType(jStartT, jO, jAi, nextTrueState);
						double nextTProb = problem.getTransProb(jStartT[0], jAi, nextTrueState) *
										   problem.getObsProb(jAi, nextTrueState, jOi);
						transitionFunction[jStartTi][jAi][nextTi] += nextTProb;
					}
				}
			}
		}
		// construct reward function (same as base reward function operating on nature's type)
		double[][] rewardFunction = new double[numJTypes][numJActions];
		for(int[] jT : new CartesianIterator(approxP.typeSS)) {
			int jTi = JointIterator.getJointIndex(approxP.typeSS, jT);
			for(int jAi = 0; jAi < numJActions; jAi++) {
				rewardFunction[jTi][jAi] = problem.getReward(jT[0], jAi);
			}
		}
		// compute starting distribution
		double[] startDistr = getStartingDist(problem.startDistr);
		approxP.setFunctions(transitionFunction, rewardFunction, startDistr);
	}

	public double[] getStartingDist(double[] baseStartDist) {
		int[] startJInds = new int[approxP.numAgents];
		for(int i=0; i<problem.numAgents; i++) {
			startJInds[i+1] = getClosestType(new NPoint(baseStartDist), i);
		}
		double[] startDistr = new double[approxP.numJTypes];
		for(int s=0; s<problem.numStates; s++) {
			startJInds[0] = s;
			int sTi = JointIterator.getJointIndex(approxP.typeSS, startJInds);
			startDistr[sTi] = baseStartDist[s];
		}
		return startDistr;
	}

	private NPoint getCentralizedBelief(int[] jStartT) {
		NPoint jointL0Belief = new NPoint(problem.numStates, 1.0);
		for(int i=0; i<problem.numAgents; i++) {
			jointL0Belief.pointMult(l0Beliefs[i].get(jStartT[i]));
		}
		jointL0Belief.scale(1.0 / jointL0Belief.l1norm());
		return jointL0Belief;
	}

	private int getNextAgentType(int agent, int curTi, int curOi, int jAi) {
		double[] curB = l0Beliefs[agent].get(curTi).getProbDist();
		NPoint nextB = problem.getNextL0Belief(agent, curB, curOi, jAi);
		// find out which type that belief belongs to
		int nextTi = getClosestType(nextB, agent);
		return nextTi;
	}
	
	/**
	 * @param jT the last joint-type
	 * @param jO the next observations seen
	 * @param jAi the last joint action taken
	 * @return the next (closest) joint-types
	 */
	private int getNextJType(int[] jT, int[] jO, int jAi, int nextTrueState) {
		int[] nextjT = new int[approxP.numAgents];
		nextjT[0] = nextTrueState;
		for(int agent=0; agent < problem.numAgents; agent++) {
			nextjT[agent+1] = getNextAgentType(agent, jT[agent+1], jO[agent], jAi);
		}
		return JointIterator.getJointIndex(approxP.typeSS, nextjT);
	}
	
	/**
	 * compute the type that has the closest euclidean distance
	 */
	public int getClosestType(NPoint b, int agent) {
		int closestT = -1;
		double bestSqDist = Double.MAX_VALUE;
		for(int t=0; t<l0Beliefs[agent].size(); t++) {
			double dist = l0Beliefs[agent].get(t).sqDistance(b);
			if(dist < bestSqDist) {
				closestT = t;
				bestSqDist = dist;
			}
		}
		return closestT;
	}
	
	public static void main(String[] args) {
		DecPOMDP tigerProb = DecPOMDPFactory.makeMATiger();
		List<Belief>[] level0Beliefs = new ArrayList[3];
		ArrayList theBs = new ArrayList<Belief>();
		theBs.add(new Belief(new NPoint(new double[]{0.5, 0.5})));
		theBs.add(new Belief(new NPoint(new double[]{0.85, 0.15})));
		theBs.add(new Belief(new NPoint(new double[]{0.15, 0.85})));
		theBs.add(new Belief(new NPoint(new double[]{0.0302, 0.9698})));
//		theBs.add(new Belief(new NPoint(new double[]{0.5, 0.5})));
		theBs.add(new Belief(new NPoint(new double[]{0.9698, 0.0302})));
		level0Beliefs[0] = theBs;
		level0Beliefs[1] = theBs;
		PPL0Approximator theApprox = new PPL0Approximator(tigerProb, level0Beliefs);
		for(int[] jT : new CartesianIterator(theApprox.approxP.typeSS)) {
			int jTi = JointIterator.getJointIndex(theApprox.approxP.typeSS, jT);
			System.out.println("State " + jTi + " is: " +
							    Arrays.toString(JointIterator.getObjects(theApprox.approxP.types, jT)));
		}
		Experiments.testDMDPVFApprox(4, 80, theApprox, 50000, 50);
	}
}
